from RumdetecFramework.BaseRumorFramework import RumorDetection
from Dataloader.twitterloader import TwitterSet, SubReader
from SentModel.Sent2Vec import *
from PropModel.SeqPropagation import GRUModel
import torch

def obtain_TheseuseLSTMRD(pretrained_file):
    lvec = BertEmb_LSTMVec("../../bert_en/", 768, 2)
    prop = GRUModel(768, 256, 1, 0.2)
    cls = nn.Linear(256, 2)
    ch = torch.load(pretrained_file)
    bvec = BertVec("../../bert_en/", bert_parallel=True)
    bvec.load_state_dict(ch['sent2vec'])
    lvec.emb.load_state_dict(bvec.bert.module.embeddings.state_dict())
    prop.load_state_dict(ch["prop_model"])
    cls.load_state_dict(ch['rdm_cls'])
    LSTMRD = RumorDetection(lvec, prop, cls)
    return LSTMRD

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set

i = 2
tr, dev, te = obtain_general_set("../data/twitter_tr%d"%i, "../data/twitter_dev%d"%i, "../data/twitter_te%d"%i)
tr.filter_short_seq(min_len=5)
tr.trim_long_seq(10)
te_event = te.data[te.data_ID[0]]['event']
print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (te_event, len(dev), len(te), len(tr)))
print("\n\n===========SubRDM Train===========\n\n")
model = obtain_TheseuseLSTMRD("../saved/G_AdverBertRD1_%s.pkl"%te_event)
model.TheseusTrain(tr, dev, te, max_epochs=100,
                log_dir="../logs/", log_suffix="_TheseusLSTMRD",
                model_file="S_TheseusLSTMRD_%s.pkl"% te_event)
